import numpy as np
import matplotlib.pyplot as plt

def solve_poisson_2d(Nx, Ny, dx, dy, f):
    """
    解二维泊松方程的函数。
    
    参数:
    Nx: int
        x方向的网格点数目。
    Ny: int
        y方向的网格点数目。
    dx: float
        x方向上的网格间距。
    dy: float
        y方向上的网格间距。
    f: numpy.ndarray
        源项矩阵，形状为 (Nx, Ny)。
    
    返回:
    u: numpy.ndarray
        数值解，即压力场。

    """
    # 网格位置
    x = np.linspace(0, Nx * dx, Nx, endpoint=False)
    y = np.linspace(0, Ny * dy, Ny, endpoint=False)
    X, Y = np.meshgrid(x, y, indexing='ij')
    
    # 构造k空间中的泊松方程系数
    kx = 2 * np.pi * np.fft.fftfreq(Nx, d=dx)
    ky = 2 * np.pi * np.fft.fftfreq(Ny, d=dy)
    lhs = -np.outer(kx**2, np.ones(Ny)) - np.outer(np.ones(Nx), ky**2)
    
    # 避免除以0
    lhs[0, 0] = 1
    
    # 计算频域中的源项
    f_hat = np.fft.fft2(f)
    
    # 解泊松方程
    u_hat = f_hat / lhs
    u = np.fft.ifft2(u_hat).real
    
    
    return u

# 主程序
if __name__ == "__main__":
    # 求解域和网格参数
    Nx = 32
    Ny = 32
    L = 1.0  # 求解域长度
    dx = L / Nx
    dy = L / Ny
    
    # 创建网格
    x = np.linspace(0, L, Nx)
    y = np.linspace(0, L, Ny)
    X, Y = np.meshgrid(x, y, indexing='ij')
    
    # 泊松方程等式右边项
    f = np.zeros((Nx,Ny))
    rsq = np.zeros((Nx,Ny))
    ft = 0
    for i in range(Nx):
        for j in range(Ny):
            rsq[i,j] = (i*dx - L/2)**2 + (j*dx - L/2)**2
            sigsq = 0.01
            f[i,j] = np.exp(-rsq[i,j] / (2*sigsq))*(rsq[i,j] - 2*sigsq)/(sigsq**2)
    
    # 调用求解函数
    u= solve_poisson_2d(Nx, Ny, dx, dy, f)

    # 检验结果，用解析及和数值解比较
    analytic = np.zeros((Nx,Ny))
    error = np.zeros((Nx,Ny)) 
    for i in range(1,Nx-1):
        for j in range(1,Ny-1):
            analytic[i,j] = np.exp(-rsq[i,j] / (2 * sigsq))
            error[i,j] = abs(analytic[i,j] - u[i,j])
    
    # # 绘图比较数值解和解析解
    # plt.figure(figsize=(12, 6))
    
    # plt.subplot(1, 2, 1)
    # plt.title("Numerical Solution")
    # plt.imshow(u, extent=[0, L, 0, L], origin='lower', cmap='viridis')
    # plt.colorbar(label='Pressure')
    
    # plt.subplot(1, 2, 2)
    # plt.title("Analytical Solution")
    # plt.imshow(analytic, extent=[0, L, 0, L], origin='lower', cmap='viridis')
    # plt.colorbar(label='Pressure')
    
    # plt.tight_layout()
    # plt.show()
    
    # 打印最大误差和均方误差
    print("Maximum Error:", np.max(error))
    print("Mean Squared Error:", np.mean(error**2))
            
    # 为了绘图，我们需要将解析解和数值解限制在相同的范围内
    # 并且移除可能由于FFT导致的高频噪声造成的误差
    u_analytic = analytic - np.mean(analytic)  # 移除解析解的均值
    u_numeric = u.real - np.mean(u.real)      # 移除数值解的均值

    # 绘制数值解和解析解
    plt.figure(figsize=(12, 6))

    # 第一个子图：数值解
    plt.subplot(1, 2, 1)
    plt.title("Numerical Solution")
    plt.imshow(u_numeric, extent=[0, L, 0, L], origin='lower', cmap='viridis')
    plt.colorbar(label='Pressure')

    # 第二个子图：解析解
    plt.subplot(1, 2, 2)
    plt.title("Analytical Solution")
    plt.imshow(u_analytic, extent=[0, L, 0, L], origin='lower', cmap='viridis')
    plt.colorbar(label='Pressure')

    # 调整子图间距
    plt.tight_layout()

    # 显示误差
    plt.figure()
    plt.title("Error between Numerical and Analytical Solutions")
    plt.imshow(error, extent=[0, L, 0, L], origin='lower', cmap='coolwarm')
    plt.colorbar(label='Error')

    # 显示所有图形
    plt.show()